/* -*- Mode: C++; c-basic-offset: 2 -*- */

#include "BcLib.h"
#include "SpecialFunctionHandler.h"

#include "llvm/Bitcode/ReaderWriter.h"

#include <vector>

//#define dprintf(args...) fprintf(stderr, args)
#define dprintf(args...)

using namespace klee;
using namespace llvm;

// Heming: original value was too big for 32 bit, so I make it smaller.
const void* FAKE_HANDLE = (const void*) 0xdeadbeef;

bool BcLib::callDLFunction(ExecutionState &state,
                           KInstruction *target,
                           Function *function,
                           std::vector< ref<Expr> > &arguments)
{
  if(function->getName() == "dlopen") {
    assert(arguments.size() == 2 && "wrong number of arguments!");
    if(isBcFile(state, arguments[0]))  {
      callDLOpen(state, target, function, arguments);
      return true;
    }
  } else if(function->getName() == "dlsym") {
    assert(arguments.size() == 2 && "wrong number of arguments!");
    if(isBcHandle(state, arguments[0])) {
      callDLSym(state, target, function, arguments);
      return true;
    }
  } else if(function->getName() == "dlclose") {
    assert(arguments.size() == 1 && "wrong number of arguments!");
    if(isBcHandle(state, arguments[0])) {
      callDLClose(state, target, function, arguments);
      return true;
    }
  } 
  
  return false;
}


void BcLib::callDLOpen(ExecutionState &state,
                       KInstruction *target,
                       llvm::Function *function,
                       std::vector< ref<Expr> > &arguments)
{
  dprintf(stderr, "BcLib::callDLOpen\n");
  // set fake handle to be the return value of the call
  const Type *resultType = target->inst->getType();
  assert(resultType != Type::getVoidTy(getGlobalContext()));

  ref<Expr> e = ConstantExpr::fromMemory((void*)&FAKE_HANDLE, 
                                         executor.getWidthForLLVMType(resultType));
  executor.bindLocal(target, state, e);
}

void BcLib::callDLSym(ExecutionState &state,
                      KInstruction *target,
                      llvm::Function *function,
                      std::vector< ref<Expr> > &arguments)
{
  dprintf(stderr, "BcLib::callDLSym\n");
  std::string symbol
    = executor.specialFunctionHandler->readStringAtAddress(state, arguments[1]);
  StringRef symbolRef(symbol);
  KModule *kmodule = executor.kmodule;
  GlobalValue *c;

  c = kmodule->module->getFunction(symbolRef);
  if(!c)
    c = kmodule->module->getNamedGlobal(symbolRef);
  if(!c)
    c = kmodule->module->getNamedAlias(symbolRef);

  assert(c&&"can't find symbol! did you run \"klee -bcload <path to bclib>?\"");

  std::map<const llvm::GlobalValue*, ref<ConstantExpr> >::iterator it
    = executor.globalAddresses.find(c);
  
  assert(it!=executor.globalAddresses.end() && "can't resolve symbol!");

  ref<Expr> e = it->second;
  const Type *resultType = target->inst->getType();
  assert(resultType != Type::getVoidTy(getGlobalContext()));
  executor.bindLocal(target, state, e);
}

void BcLib::callDLClose(ExecutionState &state,
                        KInstruction *target,
                        llvm::Function *function,
                        std::vector< ref<Expr> > &arguments)
{
  dprintf(stderr, "BcLib::callDLSym\n");
  const Type *resultType = target->inst->getType();
  assert(resultType != Type::getVoidTy(getGlobalContext()));
  ref<Expr> e = ConstantExpr::create(0, executor.getWidthForLLVMType(resultType));
  executor.bindLocal(target, state, e);
}

bool BcLib::isBcFile(ExecutionState &state, ref<Expr> path)
{
  std::string strpath 
    = executor.specialFunctionHandler->readStringAtAddress(state, path);
  
  Module *dummy = ParseBitcodeFile(MemoryBuffer::getFile(strpath), 
                                   getGlobalContext());
  bool ret = (dummy!=NULL);
  delete dummy;

  return ret;
}
 
bool BcLib::isBcHandle(ExecutionState &state, ref<Expr> handle)
{
  assert(isa<ConstantExpr>(handle) && "dl handle is not constant!");
  ref<ConstantExpr> address = cast<ConstantExpr>(handle);
  void* ptrhandle = (void*)address->getZExtValue();
  return ptrhandle == FAKE_HANDLE;
}

BcLib::BcLib(Executor &_executor)
  : executor(_executor)
{
}
